#include <iostream>
#include <cstring>
#include <chrono>
#include <glm/mat4x4.hpp>
#include <glm/gtx/transform.hpp>
#include "VK_UniformBuffer.h"
#include "VK_Context.h"
#include "VK_Image.h"
#include "VK_Texture.h"
#include "VK_DynamicState.h"
#include "VK_Pipeline.h"
#include "VK_Window.h"
#include "../Node.h"
#include "../ObjectFactory.h"

using namespace std;

#define INSTANCE_COUNT_X 15
#define INSTANCE_COUNT_Y 4

class instance: public Node
{
public:
    instance();
protected:
    RunningState run() override;
    uint32_t updateUniformBufferData(char *&data, uint32_t size);

    bool checkDevice(const VkPhysicalDeviceFeatures &feature) override
    {
        return feature.multiDrawIndirect && feature.fillModeNonSolid && feature.drawIndirectFirstInstance;
    }
protected:
    VK_Context *context = nullptr;
    VK_Pipeline *pipeline = nullptr;
};

instance::instance()
{
    nodeName = "instance";
}

RunningState instance::run()
{
    auto window = VK_Window::getInstance()->getWindow();
    auto title = getNodeTitle();
    glfwSetWindowTitle(window, title.data());

    context = createVkContext();
    context->setPhysicalDeviceFeaturesCallback(std::bind(&instance::checkDevice, this,
                                                         std::placeholders::_1));

    VkPhysicalDeviceFeatures deviceFeatures{};
    deviceFeatures.multiDrawIndirect = VK_TRUE;
    deviceFeatures.fillModeNonSolid = VK_TRUE;
    deviceFeatures.drawIndirectFirstInstance = VK_TRUE;
    context->setLogicalDeviceFeatures(deviceFeatures);

    VK_Context::VK_Config vkConfig;
    bool ok = context->initVulkanDevice(vkConfig);
    if (!ok)
        return RunningState_Failed;

    auto shaderSet = context->createShaderSet();
    shaderSet->addShader("../shader/model-instance/vert.spv", VK_SHADER_STAGE_VERTEX_BIT);
    shaderSet->addShader("../shader/model-instance/frag.spv", VK_SHADER_STAGE_FRAGMENT_BIT);

    shaderSet->appendVertexAttributeDescription(0, sizeof (float) * 3, VK_FORMAT_R32G32B32_SFLOAT, 0,
                                                0);
    shaderSet->appendVertexAttributeDescription(1, sizeof (float) * 2, VK_FORMAT_R32G32_SFLOAT,
                                                sizeof(float) * 3, 0);
    shaderSet->appendVertexAttributeDescription(2, sizeof (float) * 3, VK_FORMAT_R32G32B32_SFLOAT,
                                                sizeof(float) * 5, 0);

    shaderSet->appendVertexAttributeDescription(3, sizeof(glm::vec3), VK_FORMAT_R32G32B32_SFLOAT,
                                                0, 1);

    shaderSet->appendVertexInputBindingDescription(8 * sizeof(float), 0, VK_VERTEX_INPUT_RATE_VERTEX);
    shaderSet->appendVertexInputBindingDescription(sizeof(glm::vec3), 1, VK_VERTEX_INPUT_RATE_INSTANCE);

    VkDescriptorSetLayoutBinding uniformBinding = VK_ShaderSet::createDescriptorSetLayoutBinding(0,
                                                                                                 VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_VERTEX_BIT);
    shaderSet->addDescriptorSetLayoutBinding(uniformBinding);

    if (!shaderSet->isValid()) {
        std::cerr << "invalid shaderSet" << std::endl;
        shaderSet->release();
        context->release();
        return RunningState_Failed;
    }

    auto ubo = shaderSet->addUniformBuffer(0, sizeof(float) * 17);
    ubo->setWriteDataCallback([&](char *&data, uint32_t size)->uint32_t{return updateUniformBufferData(data, size);});

    context->initVulkanContext();

    pipeline = context->createPipeline(shaderSet);

    auto rasterCreateInfo = pipeline->getRasterizationStateCreateInfo();
    rasterCreateInfo.polygonMode = VK_POLYGON_MODE_LINE;
    pipeline->setRasterizationStateCreateInfo(rasterCreateInfo);

    pipeline->create();

    std::vector<glm::vec3> instance(INSTANCE_COUNT_X * INSTANCE_COUNT_Y);

    for (uint32_t y = 0; y < INSTANCE_COUNT_Y; y++)
    {
        for (uint32_t x = 0; x < INSTANCE_COUNT_X; x++)
        {
            if (INSTANCE_COUNT_X > 1)
            {
                instance[y * INSTANCE_COUNT_X + x].x = 7.0f * (float)x / (float)(INSTANCE_COUNT_X - 1) - 3.5f;
            }

            if (INSTANCE_COUNT_Y > 1)
            {
                instance[y * INSTANCE_COUNT_X + x].y = 5.0f * (float)y / (float)(INSTANCE_COUNT_Y - 1) - 2.5f;
            }
        }
    }

    auto buffer = context->createVertexBuffer("../model/pug.obj", true, true);
    pipeline->addRenderBuffer(buffer);

    uint32_t count = buffer->getDataCount();

    buffer = context->createInstanceBuffer(instance.size(), sizeof(glm::vec3),
                                           (char *)instance.data(), 1);
    pipeline->addRenderBuffer(buffer);

    buffer = context->createIndirectBuffer(INSTANCE_COUNT_Y, INSTANCE_COUNT_X, 0, count);
    pipeline->addRenderBuffer(buffer);

    context->createCommandBuffers();

    auto result = context->run();
    setAverageFps(context->getFps());
    context->release();
    return result;
}

uint32_t instance::updateUniformBufferData(char *&data, uint32_t size)
{
    static auto startTime = std::chrono::high_resolution_clock::now();
    auto currentTime = std::chrono::high_resolution_clock::now();
    float time = std::chrono::duration<float, std::chrono::seconds::period>
                 (currentTime - startTime).count();
    glm::mat4 model = glm::rotate(glm::mat4(1.0f), glm::radians(90.0f), glm::vec3(1.0f, 0.0f, 0.0f));
    model *= glm::rotate(glm::mat4(1.0f), time * glm::radians(30.0f), glm::vec3(0.0f, 1.0f, 0.0f));
    auto view = glm::lookAt(glm::vec3(0.0f, 4.0f, 0.0f), glm::vec3(0.0f, 0.0f, 0.0f), glm::vec3(0.0f,
                                                                                                0.0f, 1.0f));
    auto window = VK_Window::getInstance()->getWindow();
    int32_t width, height;
    glfwGetFramebufferSize(window, &width, &height);

    auto proj = glm::perspective(glm::radians(45.0f),
                                 (float)width / (float)height, 0.1f, 10.0f);
    proj[1][1] *= -1;
    model = proj * view * model;
    memcpy(data, &model[0][0], size);
    time = sin(time);
    memcpy(data + sizeof(float) * 16, (void *)&time, sizeof(float));
    return 17 * sizeof(float);
}

REGISTER_OBJECT(Node, instance)
